/*
 * Copyright (c) 2011-2022, baomidou (jobob@qq.com).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.baomidou.mybatisplus.advance.injector.processor;


import com.baomidou.mybatisplus.advance.injector.ASTBuilder;
import com.baomidou.mybatisplus.annotation.TableName;
import com.google.auto.service.AutoService;
import com.sun.source.util.Trees;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.model.JavacElements;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.Names;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import java.util.LinkedHashSet;
import java.util.Set;

/**
 * 为实体类注入相关属性，方法和mapper接口
 * 用于增强实体类功能
 * @author wanglei
 * @since 2022-03-14
 */
@SupportedAnnotationTypes("com.baomidou.mybatisplus.annotation.TableName")
@AutoService(Processor.class)
public class MybatisPlusModelInjectorAnnotationProcessor extends AbstractProcessor {

    private static Set<MybatisPlusAPTProcessor> mybatisPlusAPTProcessors = new LinkedHashSet<>();

    static{
        mybatisPlusAPTProcessors.add(new GenMapperProcessor());
        mybatisPlusAPTProcessors.add(new GenFieldsConstantProcessor());
        mybatisPlusAPTProcessors.add(new GenFieldsMethodProcessor());
        mybatisPlusAPTProcessors.add(new GenOperatorMethodProcessor());
        mybatisPlusAPTProcessors.add(new GenQueryMethodProcessor());
        mybatisPlusAPTProcessors.add(new GenImplModelExtensionMethodProcessor());
    }

    private Context context;

    private JavacElements elementUtils;

    private TreeMaker treeMaker;

    private JavacTrees javacTrees;

    private Names names;

    private Trees trees;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        trees = Trees.instance(processingEnv);
        context = ((JavacProcessingEnvironment) processingEnv).getContext();
        elementUtils = (JavacElements) processingEnv.getElementUtils();
        treeMaker = TreeMaker.instance(context);
        javacTrees = JavacTrees.instance(processingEnv);
        super.init(processingEnv);
        names = Names.instance(context);
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        //获取注解标识的类集合
        Set<? extends Element> set = roundEnv.getElementsAnnotatedWith(TableName.class);
        for (Element element : set) {
            //获取语法树的根节点
            JCTree tree = javacTrees.getTree(element);
            treeMaker.pos = tree.pos;
            //添加访问事件
            tree.accept(new TreeTranslator() {
                            @Override
                            public void visitClassDef(JCTree.JCClassDecl jcClassDecl) {
                                super.visitClassDef(jcClassDecl);
                                ASTBuilder aptBuilder = new ASTBuilder( treeMaker, elementUtils,  names,jcClassDecl);
                                for (MybatisPlusAPTProcessor mybatisPlusAPTProcessor : mybatisPlusAPTProcessors) {
                                    mybatisPlusAPTProcessor.generate(aptBuilder);
                                }
                            }
                        }
            );
        }
        return true;
    }

}
